import cv2
import os


def main(save_path='E:/yolov8/robocup/ex_train/images'):
    # 如果保存路径不存在，则创建它
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 打开摄像头
    cap = cv2.VideoCapture(0)

    # 初始化计数器
    img_count = 139

    while True:
        # 读取摄像头帧
        ret, frame = cap.read(0)
        if not ret:
            print("无法捕获图像")
            break


        # 在窗口中显示摄像头画面
        cv2.imshow('Camera', frame)

        # 按下空格键时保存当前画面帧到本地文件
        key = cv2.waitKey(1)
        if key == 32:  # 空格键的ASCII码
            img_name = os.path.join(save_path, f'image_{img_count}.jpg')
            cv2.imwrite(img_name, frame)
            print(f"已保存图像: {img_name}")
            img_count += 1

        # 按下Esc键退出循环
        elif key == 27:
            break

    # 释放摄像头并关闭窗口
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    # 修改此处的参数以调整保存路径和对比度
    main(save_path='E:/yolov8/robocup/ex_train/images')
